{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e1fcc525-8149-451e-b9f6-f1656a406d39",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "!pip install \"burr[start]\"\n",
    "!pip install opentelemetry-instrumentation-openai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f7d4516-c5e1-4889-abed-155645c41d9d",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# OpenTelemetry\n",
    "\n",
    "In this notebook we'll instrument a simple Burr app with [OpenTelemetry](www.opentelemetry.io). See the [docs](https://burr.dagworks.io/concepts/additional-visibility/) for more information."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e746862-d3b3-42f3-a3f7-0dbbdea54472",
   "metadata": {},
   "source": [
    "# Our application\n",
    "\n",
    "Let's take a look at the graph we've defined in [application.py](./application.py). As you can see, it's a simple multi-modal chatbot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1540b34e-dfd8-49f6-8d6c-5585af0596ab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 2.43.0 (0)\n",
       " -->\n",
       "<!-- Title: %3 Pages: 1 -->\n",
       "<svg width=\"740pt\" height=\"413pt\"\n",
       " viewBox=\"0.00 0.00 739.50 413.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 409)\">\n",
       "<title>%3</title>\n",
       "<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-409 735.5,-409 735.5,4 -4,4\"/>\n",
       "<!-- prompt -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>prompt</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M697,-337C697,-337 648,-337 648,-337 642,-337 636,-331 636,-325 636,-325 636,-312 636,-312 636,-306 642,-300 648,-300 648,-300 697,-300 697,-300 703,-300 709,-306 709,-312 709,-312 709,-325 709,-325 709,-331 703,-337 697,-337\"/>\n",
       "<text text-anchor=\"middle\" x=\"672.5\" y=\"-314.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">prompt</text>\n",
       "</g>\n",
       "<!-- check_safety -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>check_safety</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M675,-269C675,-269 586,-269 586,-269 580,-269 574,-263 574,-257 574,-257 574,-244 574,-244 574,-238 580,-232 586,-232 586,-232 675,-232 675,-232 681,-232 687,-238 687,-244 687,-244 687,-257 687,-257 687,-263 681,-269 675,-269\"/>\n",
       "<text text-anchor=\"middle\" x=\"630.5\" y=\"-246.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">check_safety</text>\n",
       "</g>\n",
       "<!-- prompt&#45;&gt;check_safety -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>prompt&#45;&gt;check_safety</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M661.24,-299.81C657.01,-293.15 652.09,-285.43 647.45,-278.13\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"650.11,-275.79 641.79,-269.24 644.2,-279.55 650.11,-275.79\"/>\n",
       "</g>\n",
       "<!-- input__prompt -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>input__prompt</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"731.5,-405 613.5,-405 613.5,-368 731.5,-368 731.5,-405\"/>\n",
       "<text text-anchor=\"middle\" x=\"672.5\" y=\"-382.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: prompt</text>\n",
       "</g>\n",
       "<!-- input__prompt&#45;&gt;prompt -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>input__prompt&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M672.5,-367.81C672.5,-361.6 672.5,-354.45 672.5,-347.6\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"676,-347.24 672.5,-337.24 669,-347.24 676,-347.24\"/>\n",
       "</g>\n",
       "<!-- decide_mode -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>decide_mode</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M288.5,-187C288.5,-187 198.5,-187 198.5,-187 192.5,-187 186.5,-181 186.5,-175 186.5,-175 186.5,-162 186.5,-162 186.5,-156 192.5,-150 198.5,-150 198.5,-150 288.5,-150 288.5,-150 294.5,-150 300.5,-156 300.5,-162 300.5,-162 300.5,-175 300.5,-175 300.5,-181 294.5,-187 288.5,-187\"/>\n",
       "<text text-anchor=\"middle\" x=\"243.5\" y=\"-164.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">decide_mode</text>\n",
       "</g>\n",
       "<!-- check_safety&#45;&gt;decide_mode -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>check_safety&#45;&gt;decide_mode</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M573.81,-237.78C503.93,-223.34 385.31,-198.82 310.86,-183.42\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"311.16,-179.91 300.65,-181.31 309.74,-186.77 311.16,-179.91\"/>\n",
       "<text text-anchor=\"middle\" x=\"510\" y=\"-205.8\" font-family=\"Times,serif\" font-size=\"14.00\">safe=True</text>\n",
       "</g>\n",
       "<!-- response -->\n",
       "<g id=\"node10\" class=\"node\">\n",
       "<title>response</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M561.5,-37C561.5,-37 499.5,-37 499.5,-37 493.5,-37 487.5,-31 487.5,-25 487.5,-25 487.5,-12 487.5,-12 487.5,-6 493.5,0 499.5,0 499.5,0 561.5,0 561.5,0 567.5,0 573.5,-6 573.5,-12 573.5,-12 573.5,-25 573.5,-25 573.5,-31 567.5,-37 561.5,-37\"/>\n",
       "<text text-anchor=\"middle\" x=\"530.5\" y=\"-14.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">response</text>\n",
       "</g>\n",
       "<!-- check_safety&#45;&gt;response -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>check_safety&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M640.56,-231.91C658.95,-197.19 692.95,-118.82 657.5,-68 640.8,-44.06 610.58,-31.93 583.82,-25.79\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"584.26,-22.3 573.76,-23.71 582.85,-29.16 584.26,-22.3\"/>\n",
       "</g>\n",
       "<!-- generate_image -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>generate_image</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M493.5,-105C493.5,-105 383.5,-105 383.5,-105 377.5,-105 371.5,-99 371.5,-93 371.5,-93 371.5,-80 371.5,-80 371.5,-74 377.5,-68 383.5,-68 383.5,-68 493.5,-68 493.5,-68 499.5,-68 505.5,-74 505.5,-80 505.5,-80 505.5,-93 505.5,-93 505.5,-99 499.5,-105 493.5,-105\"/>\n",
       "<text text-anchor=\"middle\" x=\"438.5\" y=\"-82.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">generate_image</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_image -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M248.76,-149.9C252.77,-139.56 259.36,-127.16 269.5,-120 301.29,-97.54 318.84,-110.82 361.33,-104.9\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"362.13,-108.31 371.41,-103.2 360.97,-101.41 362.13,-108.31\"/>\n",
       "<text text-anchor=\"middle\" x=\"352.5\" y=\"-123.8\" font-family=\"Times,serif\" font-size=\"14.00\">mode=generate_image</text>\n",
       "</g>\n",
       "<!-- generate_code -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>generate_code</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M637,-105C637,-105 536,-105 536,-105 530,-105 524,-99 524,-93 524,-93 524,-80 524,-80 524,-74 530,-68 536,-68 536,-68 637,-68 637,-68 643,-68 649,-74 649,-80 649,-80 649,-93 649,-93 649,-99 643,-105 637,-105\"/>\n",
       "<text text-anchor=\"middle\" x=\"586.5\" y=\"-82.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">generate_code</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_code -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M300.71,-160.27C339.15,-154.82 390.77,-146.33 435.5,-135 455.53,-129.93 459.83,-126.34 479.5,-120 491.87,-116.01 505.09,-111.89 517.86,-107.97\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"518.98,-111.29 527.52,-105.03 516.94,-104.6 518.98,-111.29\"/>\n",
       "<text text-anchor=\"middle\" x=\"557\" y=\"-123.8\" font-family=\"Times,serif\" font-size=\"14.00\">mode=generate_code</text>\n",
       "</g>\n",
       "<!-- answer_question -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>answer_question</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M341.5,-105C341.5,-105 225.5,-105 225.5,-105 219.5,-105 213.5,-99 213.5,-93 213.5,-93 213.5,-80 213.5,-80 213.5,-74 219.5,-68 225.5,-68 225.5,-68 341.5,-68 341.5,-68 347.5,-68 353.5,-74 353.5,-80 353.5,-80 353.5,-93 353.5,-93 353.5,-99 347.5,-105 341.5,-105\"/>\n",
       "<text text-anchor=\"middle\" x=\"283.5\" y=\"-82.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">answer_question</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;answer_question -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M186.19,-165.83C135.42,-163 66.76,-155.4 49.5,-135 45.19,-129.91 45.12,-125.03 49.5,-120 50.8,-118.51 135.83,-106.97 203.01,-98.06\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"203.75,-101.49 213.2,-96.71 202.83,-94.55 203.75,-101.49\"/>\n",
       "<text text-anchor=\"middle\" x=\"135.5\" y=\"-123.8\" font-family=\"Times,serif\" font-size=\"14.00\">mode=answer_question</text>\n",
       "</g>\n",
       "<!-- prompt_for_more -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>prompt_for_more</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M129,-105C129,-105 12,-105 12,-105 6,-105 0,-99 0,-93 0,-93 0,-80 0,-80 0,-74 6,-68 12,-68 12,-68 129,-68 129,-68 135,-68 141,-74 141,-80 141,-80 141,-93 141,-93 141,-99 135,-105 129,-105\"/>\n",
       "<text text-anchor=\"middle\" x=\"70.5\" y=\"-82.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">prompt_for_more</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;prompt_for_more -->\n",
       "<g id=\"edge11\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;prompt_for_more</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M186.14,-166.23C133.38,-163.71 60.66,-156.33 42.5,-135 36.97,-128.5 38.68,-120.77 43.27,-113.46\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"46.34,-115.2 49.71,-105.15 40.81,-110.91 46.34,-115.2\"/>\n",
       "</g>\n",
       "<!-- generate_image&#45;&gt;response -->\n",
       "<g id=\"edge12\" class=\"edge\">\n",
       "<title>generate_image&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M463.15,-67.81C473.76,-60.21 486.31,-51.2 497.68,-43.05\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"500.01,-45.68 506.1,-37.01 495.93,-39.99 500.01,-45.68\"/>\n",
       "</g>\n",
       "<!-- input__model -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>input__model</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"649.5,-187 539.5,-187 539.5,-150 649.5,-150 649.5,-187\"/>\n",
       "<text text-anchor=\"middle\" x=\"594.5\" y=\"-164.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: model</text>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_image -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M609.87,-149.76C616.31,-140.12 620.8,-128.49 613.5,-120 600.47,-104.84 548.58,-108.18 515.86,-104.71\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"515.97,-101.19 505.58,-103.24 514.98,-108.12 515.97,-101.19\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_code -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M626.11,-149.99C636.81,-141.67 644.66,-131.11 638.5,-120 636.87,-117.07 634.86,-114.36 632.58,-111.88\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"634.82,-109.18 625.06,-105.06 630.12,-114.37 634.82,-109.18\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;answer_question -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M539.29,-159.96C508.96,-154.84 471,-146.81 438.5,-135 425.23,-130.18 423.66,-125.12 410.5,-120 401.11,-116.35 383.03,-111.33 363.4,-106.34\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"364.21,-102.94 353.66,-103.89 362.51,-109.73 364.21,-102.94\"/>\n",
       "</g>\n",
       "<!-- generate_code&#45;&gt;response -->\n",
       "<g id=\"edge14\" class=\"edge\">\n",
       "<title>generate_code&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M571.49,-67.81C565.59,-60.86 558.7,-52.74 552.27,-45.16\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"554.69,-42.6 545.55,-37.24 549.35,-47.13 554.69,-42.6\"/>\n",
       "</g>\n",
       "<!-- answer_question&#45;&gt;response -->\n",
       "<g id=\"edge13\" class=\"edge\">\n",
       "<title>answer_question&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M349.04,-67.99C389.19,-57.26 439.9,-43.71 477.5,-33.66\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"478.5,-37.02 487.25,-31.06 476.69,-30.26 478.5,-37.02\"/>\n",
       "</g>\n",
       "<!-- prompt_for_more&#45;&gt;response -->\n",
       "<g id=\"edge15\" class=\"edge\">\n",
       "<title>prompt_for_more&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M141.12,-75.37C233.83,-62.07 394.53,-39.01 477.51,-27.1\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"478.08,-30.56 487.48,-25.67 477.09,-23.63 478.08,-30.56\"/>\n",
       "</g>\n",
       "<!-- response&#45;&gt;prompt -->\n",
       "<g id=\"edge16\" class=\"edge\">\n",
       "<title>response&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M573.87,-26.1C628.55,-35.66 715.5,-55.58 715.5,-85.5 715.5,-251.5 715.5,-251.5 715.5,-251.5 715.5,-266.47 707.64,-280.7 698.64,-292.09\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"695.84,-289.98 691.99,-299.85 701.15,-294.54 695.84,-289.98\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "Graph(actions=[prompt: {} -> chat_history, prompt, check_safety: prompt -> safe, decide_mode: prompt -> mode, generate_image: prompt, chat_history, mode -> response, generate_code: prompt, chat_history, mode -> response, answer_question: prompt, chat_history, mode -> response, prompt_for_more: prompt, chat_history -> response, response: response, safe, mode -> chat_history], transitions=[Transition(from_=prompt: {} -> chat_history, prompt, to=check_safety: prompt -> safe, condition=condition: default), Transition(from_=check_safety: prompt -> safe, to=decide_mode: prompt -> mode, condition=condition: safe=True), Transition(from_=check_safety: prompt -> safe, to=response: response, safe, mode -> chat_history, condition=condition: default), Transition(from_=decide_mode: prompt -> mode, to=generate_image: prompt, chat_history, mode -> response, condition=condition: mode=generate_image), Transition(from_=decide_mode: prompt -> mode, to=generate_code: prompt, chat_history, mode -> response, condition=condition: mode=generate_code), Transition(from_=decide_mode: prompt -> mode, to=answer_question: prompt, chat_history, mode -> response, condition=condition: mode=answer_question), Transition(from_=decide_mode: prompt -> mode, to=prompt_for_more: prompt, chat_history -> response, condition=condition: default), Transition(from_=generate_image: prompt, chat_history, mode -> response, to=response: response, safe, mode -> chat_history, condition=condition: default), Transition(from_=answer_question: prompt, chat_history, mode -> response, to=response: response, safe, mode -> chat_history, condition=condition: default), Transition(from_=generate_code: prompt, chat_history, mode -> response, to=response: response, safe, mode -> chat_history, condition=condition: default), Transition(from_=prompt_for_more: prompt, chat_history -> response, to=response: response, safe, mode -> chat_history, condition=condition: default), Transition(from_=response: response, safe, mode -> chat_history, to=prompt: {} -> chat_history, prompt, condition=condition: default)])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from application import graph\n",
    "graph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a07bdcf3-94f4-47d0-8ec3-00d3c7ee8c8d",
   "metadata": {},
   "source": [
    "# Building the application \n",
    "\n",
    "We're going to instrument this with OpenTelemetry -- first using Burr as the provider then using [Traceloop](traceloop.com). There are reasons to use them both, and you can use them together!\n",
    "\n",
    "First, let's use Burr to build an application (vanilla) and run it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bba29bd1-3e59-4050-821c-8da44c2c583f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from burr.core import ApplicationBuilder\n",
    "app = (\n",
    "    ApplicationBuilder()\n",
    "    .with_entrypoint(\"prompt\")\n",
    "    .with_state(chat_history=[])\n",
    "    .with_graph(graph)\n",
    "    .build()\n",
    ")\n",
    "app.run(inputs={\"prompt\":\"What is the capital of France?\"}, halt_after=[\"response\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8b7775e-012e-45da-b506-a91701d043d8",
   "metadata": {},
   "source": [
    "# Instrumenting with OpenLLMetry\n",
    "\n",
    "The [OpenLLMetry](https://github.com/traceloop/openllmetry) library provides a set of integrations that automatically instrument LLM libraries with OpenTelemetry.\n",
    "In this case we'll be using the [OpenAI instrumentor](https://pypi.org/project/opentelemetry-instrumentation-openai/).\n",
    "\n",
    "The following automatically instruments a set of attributes/spans for us:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b7a26a51-103e-4400-8c43-9ebffebd6313",
   "metadata": {},
   "outputs": [],
   "source": [
    "from opentelemetry.instrumentation.openai import OpenAIInstrumentor\n",
    "OpenAIInstrumentor().instrument()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31074710",
   "metadata": {},
   "source": [
    " Burr also provides a convenience function `init_instruments()` to instrument common 3rd party libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a5752d60",
   "metadata": {},
   "outputs": [],
   "source": [
    "from burr.integrations.opentelemetry import init_instruments\n",
    "\n",
    "init_instruments(\"openai\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bafc59c1",
   "metadata": {},
   "source": [
    "You can specify `init_all=True` instead of passing individual libraries as string. This will try to instrument currently imported libraries with the available instrumentation package you have installed (find some more [here](https://github.com/open-telemetry/opentelemetry-python-contrib/tree/main/instrumentation) and [here](https://github.com/traceloop/openllmetry/tree/main/packages)).\n",
    "\n",
    "You can get the logger to view which library is successfully instrumented."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3843f98b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "logger = logging.getLogger(\"burr.integrations.opentelemetry\")\n",
    "logger.addHandler(logging.StreamHandler())\n",
    "logger.setLevel(logging.DEBUG)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "950f701a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Instrumenting all libraries.\n",
      "`openai` is already instrumented.\n",
      "`anthropic` wasn't imported. Skipping instrumentation.\n",
      "`cohere` wasn't imported. Skipping instrumentation.\n",
      "`google.generativeai` wasn't imported. Skipping instrumentation.\n",
      "`mistralai` wasn't imported. Skipping instrumentation.\n",
      "`ollama` wasn't imported. Skipping instrumentation.\n",
      "`transformers` wasn't imported. Skipping instrumentation.\n",
      "`together` wasn't imported. Skipping instrumentation.\n",
      "`bedrock` wasn't imported. Skipping instrumentation.\n",
      "`replicate` wasn't imported. Skipping instrumentation.\n",
      "`vertexai` wasn't imported. Skipping instrumentation.\n",
      "`groq` wasn't imported. Skipping instrumentation.\n",
      "`ibm-watsonx-ai` wasn't imported. Skipping instrumentation.\n",
      "`aleph_alpha_client` wasn't imported. Skipping instrumentation.\n",
      "`pinecone` wasn't imported. Skipping instrumentation.\n",
      "`qdrant_client` wasn't imported. Skipping instrumentation.\n",
      "`chromadb` wasn't imported. Skipping instrumentation.\n",
      "`pymilvus` wasn't imported. Skipping instrumentation.\n",
      "`weaviate` wasn't imported. Skipping instrumentation.\n",
      "`lancedb` wasn't imported. Skipping instrumentation.\n",
      "`marqo` wasn't imported. Skipping instrumentation.\n",
      "`redis` wasn't imported. Skipping instrumentation.\n",
      "`langchain` wasn't imported. Skipping instrumentation.\n",
      "`llama_index` wasn't imported. Skipping instrumentation.\n",
      "`haystack` wasn't imported. Skipping instrumentation.\n",
      "`requests` is already instrumented.\n",
      "Couldn't instrument `httpx`. Package `opentelemetry-instrumentation-httpx` is missing.\n",
      "Couldn't instrument `urllib`. Package `opentelemetry-instrumentation-urllib` is missing.\n",
      "`urllib3` is already instrumented.\n"
     ]
    }
   ],
   "source": [
    "init_instruments(init_all=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c756b22-42d4-49f6-b57f-436bc7555295",
   "metadata": {},
   "source": [
    "# Tracking on Traceloop\n",
    "\n",
    "To follow this you'll need an API key for traceloop -- you can sign up for the free \n",
    "tier at [traceloop.com](traceloop.com).\n",
    "\n",
    "We can then build the application and track on traceloop. \n",
    "\n",
    "We:\n",
    "1. Initialize traceloop\n",
    "2. add the OpenTelemetryBridge hook to Burr -- this sets it up to capture Burr framework events as part of Open Telemetry\n",
    "\n",
    "Note you can use any OpenTelemetry provider -- traceloop is specifically built for this, however, and has a variety of LLM-specific integrations.\n",
    "\n",
    "Once you run this, visit [traceloop](app.traceloop.com) to see your traces!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7502ab1a-dd74-4e3d-8948-5a88ecb6ee56",
   "metadata": {},
   "outputs": [],
   "source": [
    "from traceloop.sdk import Traceloop\n",
    "import os\n",
    "from burr.integrations import opentelemetry\n",
    "from burr.core import ApplicationBuilder\n",
    "\n",
    "from traceloop.sdk import Traceloop\n",
    "\n",
    "Traceloop.init(api_key=os.environ[\"TRACELOOP_API_KEY\"])\n",
    "               \n",
    "app_otel = (\n",
    "    ApplicationBuilder()\n",
    "    .with_entrypoint(\"prompt\")\n",
    "    .with_state(chat_history=[])\n",
    "    .with_graph(graph)\n",
    "    .with_hooks(opentelemetry.OpenTelemetryBridge())\n",
    "    .build()\n",
    ")\n",
    "app.run(inputs={\"prompt\":\"What is the capital of France?\"}, halt_after=[\"response\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f04b07d-93a0-4b96-a954-d1b9d9b5b668",
   "metadata": {},
   "source": [
    "# Tracking on Burr\n",
    "\n",
    "You can also view on the Burr local tracker! This is good for local debugging and can go hand-in-hand with traceloop.\n",
    "\n",
    "To do so, all you need to do is add the `use_otel_tracing=True` parameter to `with_tracker`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0c1e26e-5b75-4bab-ad58-83aaea9f88af",
   "metadata": {},
   "outputs": [],
   "source": [
    "app = (\n",
    "    ApplicationBuilder()\n",
    "    .with_entrypoint(\"prompt\")\n",
    "    .with_state(chat_history=[])\n",
    "    .with_graph(graph)\n",
    "    .with_tracker(\n",
    "        project=\"demo_opentelemetry\",\n",
    "        use_otel_tracing=True)\n",
    "    .build()\n",
    ")\n",
    "\n",
    "app.run(inputs={\"prompt\":\"What is the capital of France?\"}, halt_after=[\"response\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c47d87b3-8d32-420b-ab75-acf2a7f6d2a9",
   "metadata": {},
   "source": [
    "# Viewing in UI\n",
    "\n",
    "The following cell gives you a link to the prior run in the UI. If you have not yet, run `burr` in your terminal to start the local tracking server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93cad8f7-9f0b-4b20-bc75-2adaca5bfde9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Markdown\n",
    "url = f\"[Link to UI](http://localhost:7241/project/demo_tracing/null/{app.uid})\"\n",
    "Markdown(url)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
